import torch
from PIL import Image
import cv2
import numpy as np

if __name__ == '__main__':
    model = torch.hub.load('../yolov5'
                           , 'yolov5s'
                           , trust_repo=True
                           , source='local')

    image_path = 'PIC1.jpg'
    img = Image.open(image_path)
    results = model(img)
    predictions = results.xyxy[0].numpy()
    original_img = cv2.imread(image_path)
    for *box, conf, cls in predictions:
        if int(cls) in [2, 3, 5, 7]:
            x_min, y_min, x_max, y_max = map(int, box)
            label = f'Vehicle {conf:.2f}'
            color = (0, 255, 0)
            thickness = 2
            cv2.rectangle(original_img, (x_min, y_min), (x_max, y_max), color, thickness)
            font = cv2.FONT_HERSHEY_SIMPLEX
            font_scale = 0.5
            font_thickness = 1
            text_size = cv2.getTextSize(label, font, font_scale, font_thickness)[0]
            cv2.rectangle(original_img, (x_min, y_min - 20), (x_min + text_size[0], y_min), color, -1)
            cv2.putText(original_img, label, (x_min, y_min - 5), font, font_scale, (255, 255, 255), font_thickness)
    cv2.imshow('Detected Vehicles', original_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
    cv2.imwrite('detected_vehicles.jpg', original_img)
